/* mvutils_mpi.c

   Implementation file for various mpi operations related to matrix-vector operations

   Basile Clout, September 2007

   From M. J. Quinn
*/


#include "mvutils_mpi.h"
#include <string.h>


/* Get the size of a MPI data type */
int sizeMPI(MPI_Datatype dt){
  if(dt == MPI_BYTE)
    return sizeof(char);
  if(dt == MPI_DOUBLE)
    return sizeof(double);
  if(dt == MPI_FLOAT)
    return sizeof(float);
  if(dt == MPI_INT)
    return sizeof(int);
  printf("Error: Unrecognized datatype!\n"); fflush(stdout);
  MPI_Abort(MPI_COMM_WORLD, 1);
}


/* Read matrix by column */
int mread(char *fn, double ***lmat, double **lar, int *m, int *n, MPI_Comm com, int rm){

  double *buffer;
  int i,j;
  int r, p;
  int dummy, type;
  FILE *matS;
  int lc;			/* local number of columns */
  int nb;
  double *lptr, *rptr;
  double **tmat, *tar;
  
  int *scatter_count;
  int *scatter_displ;
    
  MPI_Comm_size(com, &p);
  MPI_Comm_rank(com, &r);

  if (r==rm){
    
    if(!(matS = fopen(fn, "r"))){
      printf("Impossible to open matrix file!\n");
      MPI_Abort(com, 1);
    }

    nb = fscanf(matS, "%d", m);
    nb += fscanf(matS, "%d", n);
    nb += fscanf(matS, "%d", &dummy);
    nb += fscanf(matS, "%d", &type);
    if(nb!=4){
      printf("Error: when reading info from matrix\n");fflush(stdout);
      MPI_Abort(com, 1);
    }
    
    if (!type)
      printf("Warning matrix: in this implementation, all int are cast to double, and all operations involve double only.\n");fflush(stdout);
  }

  
  MPI_Bcast(m, 1, MPI_INT, rm, com);
  MPI_Bcast(n, 1, MPI_INT, rm, com);
  
  /* Malloc the memory needed by each process, having lc columns of *m rows,
     containing elements of size elsize (MPI_INT, MPI_DOUBLE, MPI_FLOAT, ...)
     Special kind of two-dimensional arrays: it is continuous (no ) between each row in memory like in gcc
     Useful when we want to know exactly what we send in a MPI message!
  */


  lc = BLOCK_SIZE(r, p, (*n));
  tar = (double *) malloc(*m * lc * sizeof(double)); /* Allocate complete array for a local process*/
  if (!tar){
    printf("Error from %d: Impossible to allocate memory for local array\n", r);
    MPI_Abort(com, 1);
  }
  
  tmat = (double **) malloc(*m * sizeof(double *)); /* Simple array pointing to the correct lar elements
						       Therefore equivalent to a 2D array*/
  if (!tmat){
    printf("Error: Impossible to allocate lmat memory for %d", r);
    MPI_Abort(com, 1);
  }

  for (i=0;i<*m;i++){
    // printf("Im here! from %d\n", r);fflush(stdout);
    tmat[i] = tar + i*lc;
    //printf("lamt[i] = %.2f from %d\n", tmat[i], r );
  }


  if (r == rm){
    buffer = (double *) malloc (*n * sizeof(double)); /* Complete buffer for one line, used for the scatter */
    if(!buffer){
      printf("Impossible to allocate buffer memory for master %d", r);
      MPI_Abort(com, 1);
    }
  }

  
  /* Fill in the structures for the scatter */
  scatter_count = (int *) malloc(p * sizeof(int));
  scatter_displ = (int *) malloc(p*sizeof(int));
  if (!scatter_count || !scatter_displ){
    printf("Error: Impossible to allocate scatter memory for node %d", r);
    MPI_Abort(com, 1);
  }

  scatter_count[0] = BLOCK_SIZE(0, p, *n);
  scatter_displ[0] = 0;
  for(i=1;i<p;i++){
    scatter_displ[i] = scatter_displ[i-1] + scatter_count[i-1];
    scatter_count[i] = BLOCK_SIZE(i, p, *n);
  }


  /* Actual reading and scattering */
  for(i=0;i<*m;i++){
      
    if(r == rm){
      /* Fill in one line of the matrix in buffer */
      for(j=0;j<*n;j++){
	nb = fscanf(matS, "%lf", &buffer[j]);
	if (nb!=1){
	  printf("Error when reading matrix file by master %d\n", r);
	  MPI_Abort(com, 1);
	}
      }
    }
    
    /* Scatter the buffer */
    MPI_Scatterv(buffer, scatter_count, scatter_displ, MPI_DOUBLE, tar + i*lc, lc, MPI_DOUBLE, rm, com);

  }

  *lar = tar;
  *lmat = tmat;
  
  free(scatter_count);
  free(scatter_displ);
  if (r == rm)
    free(buffer);


  return 1;
}


int mprint(double **mat, int m, int n, MPI_Comm com, int rm){

  int r, p;
  int *gather_count;
  int *gather_displ;
  int i, j;
  double *buffer;


  MPI_Comm_rank(com, &r);
  MPI_Comm_size(com, &p);

  /* Create the arrays for the gatherv */
  gather_count = (int *) malloc(p*sizeof(int));
  gather_displ = (int *) malloc(p*sizeof(int));
  if (!gather_count || !gather_displ){
    printf("Error: Impossible to allocate gather memory for node %d", r);
    MPI_Abort(com, 1);
  }

  gather_count[0] = BLOCK_SIZE(0, p, n);
  gather_displ[0] = 0;
  for (i=1;i<p;i++){
    gather_displ[i] = gather_displ[i-1] + gather_count[i-1];
    gather_count[i] = BLOCK_SIZE(i, p, n);
  }

  /* Allocate a buffer for one line of the matrix */
  if (r==rm){
    buffer = (double *) malloc(n*sizeof(double));
    if (!buffer){
      printf("Impossible to malloc the printing buffer by master %d\n", r);
      MPI_Abort(com, 1);
    }
  }

  for(i=0;i<m;i++){

    /* Gather */
    MPI_Gatherv (mat[i], BLOCK_SIZE(r, p, n), MPI_DOUBLE, buffer, gather_count, gather_displ, MPI_DOUBLE, rm, com);

    /* Printing by master */
    if (r==rm){
      for(j=0;j<n;j++)
	printf("%.2f ", buffer[j]);
      putchar('\n');
    }
    fflush(stdout);
  }

  free(gather_displ);
  free(gather_count);
  if(r==rm){
    free(buffer);
    putchar('\n');
  }
    
  
  return 1;
}


int vread(char *fn, double **subvec, int *n, MPI_Comm com, int rm){

  int r, p;
  int i, j;
  FILE *vecS;
  int mv, nv; 
  int dummy;
  int type;
  int nb=0;
  int lc=0;
  int lcv;
  double *subvec_master;
  double *tmpsubvec;
  MPI_Status status;
  
  MPI_Comm_size(com, &p);
  MPI_Comm_rank(com, &r);

  if (r==rm){

    if(!(vecS = fopen(fn, "r"))){
	printf("Impossible to open vector file %s.\n", fn);
	MPI_Abort(com, 1);
      }
          
      nb = fscanf(vecS, "%d", &mv);
      nb += fscanf(vecS, "%d", &nv);
      nb += fscanf(vecS, "%d", &dummy);
      nb += fscanf(vecS, "%d", &type);
      
      if(nb!=4){
	printf("Error: when reading info from vector\n");fflush(stdout);
	MPI_Abort(com, 1);
      }
      
      if (!type){
	printf("Warning vector: in this implementation, all int are cast to double, and all operations involve double only.\n");
	fflush(stdout);
      }

      if(mv == 1 && nv == 1){
	printf("Error: One of the dimension of the vector matrix must be 1!\n");
	fflush(stdout);
	MPI_Abort(com, 1);
      }
      
      *n = (mv == 1) ? nv : mv ; 
  }

  MPI_Bcast(n, 1, MPI_INT, rm, com);

  lc = BLOCK_SIZE(r, p, *n);
  *subvec = (double *) malloc(lc*sizeof(double));

  if(r==rm){
    
    subvec_master = (double *) malloc(*n*sizeof(double));

    /* Master send chunks of the vector to the other nodes */
    for(i=0;i<p;i++){

      lcv = BLOCK_SIZE(i, p, *n);
      tmpsubvec = (double *) malloc(lcv*sizeof(double));
      if(!tmpsubvec){
	printf("Error: Impossible to allocate temp buffer before sending in vread by %d.\n", r);
	MPI_Abort(com, 1);
      }

      for(j=0;j<lcv;j++){
	nb = fscanf(vecS, "%lf", &tmpsubvec[j]);
	if(nb!=1){
	  printf("Error: When reading vector by %d\n", r);
	  MPI_Abort(com, 1);
	}
      }

      if(i==r){			/* If master node, copy the buffer in a temp buf */
	memcpy(subvec_master, tmpsubvec, *n*sizeof(double));
      }
      else
	MPI_Send(tmpsubvec, lcv, MPI_DOUBLE, i, 1, com);
    }

    /* Reassign the the temp buffer external master buffer */
    *subvec = subvec_master;
    
  }else
    MPI_Recv(*subvec, lc, MPI_DOUBLE, rm, 1, com, &status);

  return 1;
}


  
int vprint(double *v, int n, MPI_Comm com, int rm){

  int i, r, p;
  MPI_Status status;
  int *gather_count;
  int *gather_displ;
  int lc = 0;
  double *buffer;

  MPI_Comm_size(com, &p);
  MPI_Comm_rank(com, &r);

  /* Can use MPI_Send-Recv like in the book
     Lets try a Gatherv!
     */

  if(r==rm){
    buffer = (double *) malloc(n*sizeof(double));
    if(!buffer){
      printf("Impossible to allocate memory for printing vector by %d\n", r);
      MPI_Abort(com, 1);
    }
  }
  
  gather_count = (int *) malloc(p*sizeof(int));
  gather_displ = (int *) malloc(p*sizeof(int));

  gather_count[0] = BLOCK_SIZE(0, p, n);
  gather_displ[0] = 0;
  
  for(i=1;i<p;i++){
    gather_count[i] = BLOCK_SIZE(i, p, n);
    gather_displ[i] = gather_displ[i-1] + gather_count[i-1];
    //    printf("gather tabs i = %d: cnt %d disp %d\n", i, gather_count[i], gather_displ[i]);
  }

  lc = BLOCK_SIZE(r, p, n);
  MPI_Gatherv(v, lc, MPI_DOUBLE, buffer, gather_count, gather_displ, MPI_DOUBLE, rm, com);

  /* Print the vector */
  if (PRINT){
    if(r==rm){
      for(i=0;i<n;i++)
	printf("%.2f ", buffer[i]);
      putchar('\n');
    }
  }
  
  return 1;
}

      
    

  
